import cv2
import numpy as np
import torch
from torchvision import transforms
from model import FCN8s
import utils


def inference_single_image(model, weight_path, image_path, mean, std):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    # 加载模型权重
    checkpoint = torch.load(weight_path, map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])
    model = model.to(device)
    model.eval()

    # 推理图片数据处理

    img = cv2.imread(image_path)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img_resized = cv2.resize(img_rgb, (224, 224), interpolation=cv2.INTER_LINEAR)

    # transform 和训练一致
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean,std)
    ])
    img_tensor = transform(img_resized).unsqueeze(0).to(device)  # 1x3xHxW 增加一个batch通道

    # 推理
    with torch.no_grad():
        output = model(img_tensor) # 1x2xHxW
        # 获取概率最大的分类，也就是维度为1的概率向量，并且把batch维度删掉(因为这个就是1)，转成cpu是为了方便numpy处理
        pred = torch.argmax(output, dim=1).squeeze(0).cpu().numpy()

        # 将输出结果(原本是二值掩码，现在转成标准灰度图)
        pred_mask = (pred * 255).astype(np.uint8)

        # 将掩码大小转成与图像一致大小
        pred_mask_resized = cv2.resize(pred_mask, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_NEAREST)

        # 抠图(与掩码图进行操作)，扣出目标，背景为黑色
        foreground = cv2.bitwise_and(img_rgb, img_rgb, mask=pred_mask_resized)
        return foreground, pred_mask_resized



if __name__ == "__main__":

    weight_path = "../checkpoints/best_model.pth"
    image_path = "./1.png"

    mean, std = utils.calc_mean_std_func(img_dir="../Portrait-dataset-2000/dataset/training")

    # 初始化模型
    model = FCN8s(num_classes=2)

    # 调用推理函数
    foreground, background = inference_single_image(model, weight_path, image_path, mean, std)

    # 保存结果
    cv2.imwrite("background.png", background)
    cv2.imwrite("foreground.png", cv2.cvtColor(foreground, cv2.COLOR_RGB2BGR))
    print("推理完成，结果已保存！")





